import matplotlib.pyplot as plt
import numpy as np
import torch
import pandas as pd
import torchvision.datasets
import torchvision.transforms
from PIL import Image
from p_tqdm import p_map
from tqdm import tqdm
import scipy.io as scio
# from superpixel import star_seg_M1 as star_seg
# import utils.metric
import csv
import metric
from statistics import mean
# import tensorflow as tf
import cv2
import math
import scipy.io as io
import sys
import shutil

from demo.superpixel.superpixel import SegmentorBuilder

path_cityscapes = f"dataset/"

test_name = "sh"
n_seg = 100
save_fmt = ".csv"
model = None
use_gpu = True
asdasd = 0

def modify(path):
    path = path.replace("test", "result/" + test_name).replace(".jpg", save_fmt).replace(".png", save_fmt)
    path = path.replace("val", "result/" + test_name)
    path = path.replace("JPEGImages", "result/" + test_name)
    path = path.replace("train", "result/" + test_name)
    path = path.replace("vis", "result/" + test_name)
    return path


img_cache = None
path_cache = None
result_cache = None

settings = [1, 1., 0.25]

def get_segmentor():
    seg_builder = SegmentorBuilder(d=9, objective="SE_M").convertLAB().set_gpu().constrain_neighbor_8connect() \
        .normalize_hyper_fea().iterative_mean_shift()  # .fuse_full_graph()
    return seg_builder.segmentor

def procc(path, save="csv", img=None, resize=(1, 1)):
    global save_fmt
    global img_cache
    global result_cache
    global path_cache
    if path is not None and path_cache == path and result_cache is not None:
        return result_cache
    path_cache = path
    if img is None:
        img = Image.open(path).convert('RGB')
        if resize != (1, 1):
            img = img.resize((int(img.width * resize[0]), int(img.height * resize[1])))
    # img2 = segmentation.slic(img,n_segments=n_seg,   start_label=0)
    # import mmseg.apis
    # img2 = mmseg.apis.inference_segmentor(model, path)
    # img2 = img2[0]
    arrimg = np.array(img)
    if img_cache is not None and img_cache.shape == arrimg.shape and (img_cache == arrimg).all():
        return result_cache
    img_cache = arrimg

    save_fmt = ".csv" if save == "csv" else ".png"
    ret = modify(path)
    # ret = f'/home/hujin/superpixel-benchmark/dataset/bsds/gt/'
    # print(ret)
    # img2 = star_seg.superpixel(img, target_num=n_seg, force_connectivity=True,  # r2=11, k2=60,
    #                            lambda_aggregation=settings[0], lambda_diffusion=settings[1], t=settings[2],
    #                            use_lab=True,
    #                            t_decay=1.0, lambda_pre=1.0, use_featuremap=False, scaling_bound=math.inf,
    #                            interseg=False)
    img2 = get_segmentor().run(arrimg, target_size=n_seg).cpu().numpy()


    # import superpixel.SH.src.boruvkasupix
    # img2 = superpixel.SH.src.boruvkasupix.supix(img, n_seg)
    # print(img2.max())
    # from skimage.segmentation._slic import _enforce_label_connectivity_cython
    # img2 = _enforce_label_connectivity_cython(img2.reshape(1,img2.shape[0] , img2.shape[1]).astype(np.int64),
    #                                           min(img2.shape[0] * img2.shape[1] // (20 * n_seg), 0),
    #                                           img2.shape[0] * img2.shape[1] //2, start_label=0)[0]
    #
    # print(img2.max())
    # return img2
    try:
        if os.path.exists(ret):
            os.remove(ret)
        if save == "csv":
            np.savetxt(ret, img2, fmt='%d', delimiter=",")
        else:
            Image.fromarray(img2.astype(np.int16)).save(ret)
    except:
        try:
            os.makedirs(os.path.dirname(ret))
        except:
            pass

        if save == "csv":
            np.savetxt(ret, img2, fmt='%d', delimiter=",")
        else:
            Image.fromarray(img2.astype(np.int16)).save(ret)
        # scipy.io.savemat(ret, img2, delimiter=",")
        # Image.fromarray(img2.astype(np.int16)).save(ret)
    # del img
    result_cache = img2
    return img2


def process(split='val', mode='coarse', re_process=False):
    path_superpixel = path_cityscapes + "superpixel_" + split + mode
    if not re_process:
        dataset = torchvision.datasets.Cityscapes(path_cityscapes, split=split, mode=mode,
                                                  target_type='semantic')
        return dataset, [modify(x) for x in dataset.images]
    else:
        # print("Preprocessing superpixel: {} {}".format(split, mode))
        #
        dataset = torchvision.datasets.Cityscapes(path_cityscapes, split=split, mode=mode,
                                                  target_type='semantic')
        # total_cnt = len(dataset)
        # dataset = [data[0] for data in dataset]
        results = p_map(procc, dataset.images, num_cpus=8)
        # with Pool(processes=2) as p:
        #     max_ = 30
        #     with tqdm(total=max_) as pbar:
        #         for i, _ in tqdm(enumerate(p.imap_unordered(_foo, range(0, max_)))):
        #             pbar.update()
        # with mp.Pool(4) as p:
        #     results = list(tqdm.tqdm(p.imap(procc, dataset), total=total_cnt))

        # np.savez_compressed(path_superpixel, np.array(results))
        # dataset = torchvision.datasets.Cityscapes(path_cityscapes, split=split, mode=mode,
        #                      target_type='semantic')
        return dataset, results


def proc_cache(path, save="img", img=None):
    global save_fmt
    save_fmt = ".csv" if save == "csv" else ".png"
    file = path
    if save == "csv":
        return np.fromfile(file, dtype=np.int, sep=",")
    else:
        return np.array(Image.open(file), dtype=np.int)


class CityscapesWrapped(torchvision.datasets.Cityscapes):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        if not os.path.exists(modify(self.images[0])):
            print("init slic...")
            p_map(procc, self.images, num_cpus=6)

    def __getitem__(self, index: int):
        i, s = super().__getitem__(index)
        return i, s, Image.open(modify(self.images[index]))


def process_sbsd():
    file = f"../../research/dataset/BSR/BSDS500/data/images/test/"
    imgs = [file + i for i in os.listdir(file)[:-1]]

    p_map(procc, imgs, num_cpus=6)


def file2array(path, delimiter=' '):
    fp = open(path, 'r', encoding='utf-8')
    string = fp.read()
    fp.close()
    row_list = string.splitlines()
    data_list = [[np.int32(i) for i in row.strip().split(delimiter)] for row in row_list]
    return np.array(data_list)


import random
import json

log_file = None
import os

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"  # （保证程序cuda序号与实际cuda序号对应）


# os.environ['CUDA_VISIBLE_DEVICES'] = "1"  # （代表仅使用第0，1号GPU）

def test_cityscapes_parallel(n_cpu):
    path_cityscapes = f"/home/hujin/src/Cityscapes/"
    # path_cityscapes = f"../../research/dataset/CityScapes/"
    split = 'val'
    mode = 'fine'

    class p_cs(torchvision.datasets.Cityscapes):
        def __getitem__(self, idx):
            return self.images[idx], super().__getitem__(idx)

    dataset = p_cs(path_cityscapes, split=split, mode=mode,
                   target_type='semantic',
                   transform=torchvision.transforms.Resize((512, 1024)),
                   target_transform=torchvision.transforms.Resize((512, 1024)))
    global use_gpu
    acc = 0
    i = 0
    use_gpu = False

    def proc(x):
        path, pp = x
        img, label = pp
        procc(path, save="img", img=img)

    p_map(proc, dataset, num_cpus=n_cpu)
    use_gpu = True
    dev = "cuda" if use_gpu and torch.cuda.is_available() else "cpu"
    for imgpath, (img, label) in tqdm(dataset):
        imgt = proc_cache(imgpath, save="img", img=img)
        acc += metric.calc_ASA(torch.LongTensor(imgt).to(dev), torch.LongTensor(np.array(label)).to(dev)).cpu()
        i += 1
    print(acc / len(dataset), file=log_file, flush=True)


def test_cityscapes_small():
    path_cityscapes = f"/home/hujin/src/Cityscapes/"
    split = 'val'
    mode = 'fine'
    tr = torchvision.transforms.Compose([torchvision.transforms.Resize((512, 1024)),
                                         torchvision.transforms.CenterCrop((256, 256))])
    dataset = torchvision.datasets.Cityscapes(path_cityscapes, split=split, mode=mode,
                                              target_type='semantic',
                                              transform=tr,
                                              target_transform=tr)

    acc = 0
    i = 0
    for img, label in tqdm(dataset):
        imgt = proc_cache(dataset.images[i], save="img", img=img)
        acc += metric.calc_ASA(torch.LongTensor(imgt), torch.LongTensor(np.array(label)))
        i += 1
    print(acc / len(dataset), file=log_file, flush=True)


# def test_deeplab():
#     global model
#     from mmseg.apis import init_segmentor
#     file_ckpt = f"D:/research/deeplabv3plus_r101-d8_512x512_80k_ade20k_20200615_014139-d5730af7.pth"
#     file_seg = f"../mmsegmentation/configs/deeplabv3plus/deeplabv3plus_r101-d8_512x512_80k_ade20k.py"
#     model = init_segmentor(file_seg, file_ckpt)
#     # model = torchvision.models.segmentation.deeplabv3_resnet101(pretrained=True, progress=True)
#     model.eval()
#     model.cuda()
#
#     file = f"../../research/dataset/BSR/BSDS500/data/images/train/"
#     imgs = [file + i for i in os.listdir(file)[:-1]]
#     for i in tqdm(imgs):
#         procc(i)


def eval_bsds():
    file_csvgt = "123.csv"

    dev = "cuda" if use_gpu and torch.cuda.is_available() else "cpu"
    list_seg = os.listdir(file_csvgt)


dataset = ""

train = []
from shutil import copyfile


def get_pvoc():
    dir = pd.read_csv('/home/hujin/SI/SI/SI/pvoc/voc/VOC2012/ImageSets/Segmentation/val.txt', header=None)
    dir = dir.iloc[:, 0]
    filec = '/home/hujin/SI/SI/pvoc/voc/VOC2012/JPEGImages/'
    filet = '/home/hujin/SI/SI/pvoc/voc/VOC2012/SegmentationObject/'
    images = []
    for i in dir:
        # copyfile(filec+i+'.jpg', '/home/hujin/SI/SI/pascal/test/'+i+'.jpg')
        # copyfile(filet+i+'.png', '/home/hujin/SI/SI/pascal/gt/'+i+'.png')
        train.append([filec + i + '.jpg', filet + i + '.png'])


def get_pascon():
    dir = pd.read_csv('/home/hujin/SI/SI/pascon/VOCdevkit/VOC2010/ImageSets/Segmentation/val.txt', header=None)
    dir = dir.iloc[:, 0]
    filec = '/home/hujin/SI/SI/pascon/VOCdevkit/VOC2010/JPEGImages/'
    filet = '/home/hujin/SI/SI/pascon/VOCdevkit/VOC2010/SegmentationObject/'
    filek = '/home/hujin/SI/SI/pascon/trainval/'
    for i in dir:
        if i[3] == '7':
            continue
        try:
            gt = io.loadmat(filek + i + '.mat')
        except IOError as e:
            continue
        train.append([filec + i + '.jpg', filek + i + '.mat'])


def get_mscoco():
    picdir = f'/home/hujin/SI/SI/mscoco/val2017/'
    gtdir = f'/home/hujin/SI/SI/mscoco/annotations/stuff_val2017_pixelmaps/'
    filelist = [i for i in os.listdir(picdir)]
    for i in filelist:
        train.append([picdir + i, gtdir + i[:-4] + '.png'])


def get_city():
    path_cityscapes = f"/home/hujin/SI/SI/Cityscapes/small/"
    split = 'val'
    mode = 'fine'
    dataset = torchvision.datasets.Cityscapes(path_cityscapes, split=split, mode=mode,
                                              target_type='semantic',
                                              transform=torchvision.transforms.Resize((512, 1024)),
                                              target_transform=torchvision.transforms.Resize((512, 1024)))
    # path = '/home/hujin/SI/SI/Cityscapes/gtFine/test/'
    global train

    for i in range(len(dataset.images)):
        img = dataset.images[i]
        tgt = dataset.targets[i]
        train.append([img, tgt[0]])


def get_ade():
    file = f'/home/hujin/SI/SI/ade/images/ADE/training/'
    filelist = [i for i in os.listdir(file)[:-1]]
    for i in filelist:
        pos = file + i + '/'
        flist = [j for j in os.listdir(pos)[:-1]]
        for j in flist:
            nepos = pos + j + '/'
            piclist = [k for k in os.listdir(nepos)[:-1]]
            for k in piclist:
                if k[-3:] == 'jpg' and random.randint(0, 99) <= 2:
                    train.append([nepos + k, nepos + k[:-4] + '_seg.png'])


def get_sift():
    picdir = f'/home/hujin/SI/SI/sift/sift/Images/images/'
    gtdir = f'/home/hujin/SI/SI/sift/sift/SemanticLabels/gt/'
    filelist = [i for i in os.listdir(picdir)]
    for i in filelist:
        train.append([picdir + i, gtdir + i[:-4] + '.mat'])


def get_stanf():
    picdir = f'/home/hujin/SI/SI/stab/stan/images/'
    gtdir = f'/home/hujin/SI/SI/stab/stan/labels/'
    filelist = [i for i in os.listdir(picdir)]
    for i in filelist:
        train.append([picdir + i, gtdir + i[:-4] + ".layers.txt"])


def get_bsds(ds="test"):
    file = "/home/hujin/SI/SI/bsd/BSR/BSDS500/data/images/{}/".format(ds)
    file1 = "/home/hujin/SI/SI/bsd/BSR/BSDS500/data/csv_gt/{}/".format(ds)
    imgs = [i for i in os.listdir(file1)]
    imgids = set([i.split(".")[0] for i in os.listdir(file)])
    for i in imgs:
        imgname = i.split("-")[0]
        if not i.endswith(".csv") or not imgids.__contains__(imgname):
            continue
        train.append([file + "{}.jpg".format(imgname), file1 + i])
    train.sort()


def get_kitti():
    picdir = '/home/hujin/SI/SI/kit/training/image_2/'
    gtdir = '/home/hujin/SI/SI/kit/training/semantic_rgb/'
    filelist = [i for i in os.listdir(picdir)]
    for i in filelist:
        train.append([picdir + i, gtdir + i])


def get_nyu():
    gtdir = f'/home/hujin/SI/SI/nyudv2/csv_groundTruth/test/'
    file = f'/home/hujin/SI/SI/nyudv2/test/'
    filelist = [i for i in os.listdir(file)]
    for i in filelist:
        train.append([file + i, gtdir + i[:-4] + ".csv"])


def get_sun():
    dirs = '/home/hujin/SI/SI/sun/SUNRGBD/kv1/b3dodata/'
    images = []
    filelist = [i for i in os.listdir(dirs)]
    for i in filelist:
        ndir = dirs + i
        picdir = ndir + '/image/' + i + '.jpg'
        gtdir = ndir + '/seg.mat'
        images.append([picdir, gtdir])
    dirs = '/home/hujin/SI/SI/sun/SUNRGBD/kv1/NYUdata/'
    filelist = [i for i in os.listdir(dirs)]
    for i in filelist:
        ndir = dirs + i
        picdir = ndir + '/image/' + i + '.jpg'
        gtdir = ndir + '/seg.mat'
        images.append([picdir, gtdir])

    dirs = '/home/hujin/SI/SI/sun/SUNRGBD/kv2/kinect2data/'
    filelist = [i for i in os.listdir(dirs)]
    for i in filelist:
        if i[0] == '.':
            continue
        ndir = dirs + i
        nlst = [j for j in os.listdir(str(ndir) + '/image/')]
        picdir = ndir + '/image/'
        for j in nlst:
            picdir = picdir + j
            break
        gtdir = ndir + '/seg.mat'
        images.append([picdir, gtdir])

    dirs = '/home/hujin/SI/SI/sun/SUNRGBD/kv2/align_kv2/'
    filelist = [i for i in os.listdir(dirs)]
    for i in filelist:
        if i[0] == '.':
            continue
        ndir = dirs + i
        nlst = [j for j in os.listdir(str(ndir) + '/image/')]
        picdir = ndir + '/image/'
        for j in nlst:
            picdir = picdir + j
            break
        gtdir = ndir + '/seg.mat'
        images.append([picdir, gtdir])

    for i in range(0, 300):
        train.append(random.choice(images))


def get_rgbd():
    dirs = '/home/hujin/SI/SI/rgbd/rgbd-dataset/'
    images = []
    filelist = [i for i in os.listdir(dirs)]
    for i in filelist:
        ndirs = dirs + str(i)
        filelist1 = [j for j in os.listdir(ndirs)]
        for j in filelist1:
            ndirs2 = ndirs + '/' + str(j)
            filelist2 = [k for k in os.listdir(ndirs2)]
            for k in filelist2:
                if k.find("mask") == -1 and k.find("depth") == -1 and k.find("txt") == -1:
                    fin = ndirs2 + '/' + k
                    images.append(fin)

    for i in range(0, 5000):
        sec = random.choice(images)
        gt = sec[:-8] + 'mask' + sec[-8:]
        try:
            img = Image.open(gt)
        except IOError as e:
            i -= 1
            continue
        train.append([sec, gt])


def process_dataset():
    pri = []
    vi = []
    gce = []
    asa = []
    bd1 = []
    bd2 = []
    mx_seg = []

    for i in tqdm(train):
        # print(i)
        pred = procc(i[0], resize=(x, y))
        gt = pred
        if i[1].endswith(".png"):
            img = Image.open(i[1])
            gt = np.array(img)
            if len(gt.shape) == 3:
                gt = (gt[:, :, 0] << 16) + (gt[:, :, 1] << 8) + gt[:, :, 2]
        elif i[1].endswith(".mat"):
            gt = io.loadmat(i[1])
            if dataset == "pascon":
                gt = gt['LabelMap']
            elif dataset == "sift":
                gt = gt['S']
            elif dataset == "sun":
                gt = gt['seglabel']
        elif i[1].endswith(".txt"):
            gt = file2array(i[1])
        elif i[1].endswith(".csv"):
            gt = np.loadtxt(i[1], dtype=np.int, delimiter=',')

        if x != 1.0 or y != 1.0:
            img = Image.fromarray(gt.astype('uint8')).convert()
            img = img.resize((int(img.width * x), int(img.height * y)))
            gt = np.array(img)

        void_xx = None
        if dataset == "pvoc":
            void_xx = 255
        p, v, g, a, b1, b2, c = metric.evaluate(pred, gt, void_xx)
        if np.isnan(b2):
            continue
        # print(p, v, g, a, b1, b2, c)
        pri.append(p)
        vi.append(v)
        gce.append(g)
        asa.append(a)
        bd1.append(b1)
        bd2.append(b2)
        mx_seg.append(c)
    with open('output/' + exp_name + dataset + '_result.csv', 'a') as f:
        writer = csv.writer(f)
        writer.writerow(settings)
        writer.writerow([n_seg, mean(pri), mean(vi), mean(gce), mean(asa), mean(bd1), mean(bd2), mean(mx_seg)])
    global path_cache
    path_cache = None

ava_dataset = ["pvoc", "pascon", "mscoco", "city", "ade", "sift", "stanf", "bsds", "bsds_val", "bsds_train", "kitti",
               "nyu", "sun", "rgbd"]


def get_dataset():
    picdir = f'/home/hujin/superpixel-benchmark/dataset/bsds/test/'
    gtdir = f'/home/hujin/superpixel-benchmark/dataset/bsds/gt/'
    filelist = [i for i in os.listdir(picdir)]
    for i in filelist:
        train.append([picdir + i, gtdir + i[:-4] + '.jpg'])


def init_dataset():
    # get_dataset()
    # return
    if dataset == "pvoc":
        get_pvoc()
    if dataset == "pascon":
        get_pascon()
    if dataset == "mscoco":
        get_mscoco()
    if dataset == "city":
        get_city()
    if dataset == "ade":
        get_ade()
    if dataset == "sift":
        get_sift()
    if dataset == "stanf":
        get_stanf()
    if dataset == "bsds":
        get_bsds("test")
    if dataset == "bsds_val":
        get_bsds("val")
    if dataset == "bsds_train":
        get_bsds("train")
    if dataset == "kitti":
        get_kitti()
    if dataset == "nyu":
        get_nyu()
    if dataset == "sun":
        get_sun()
    if dataset == "rgbd":
        get_rgbd()
    # return
    # global train
    # train = random.sample(train, 2000)
    # with open('/home/hujin/SI/SI/dataset_list.txt','w') as f:
    #     m = 0
    #     for i in train:
    #         print(m)
    #         img = Image.open(i[0])
    #         img = img.resize((int(img.width*x),int(img.height*y)))
    #         name = (i[0].split("/")[-1])
    #         img.save("/home/hujin/superpixel-benchmark/dataset/pascal/test/"+str(m)+".png")
    #         # img.save("/home/hujin/SI/SI/trainset/rgbd/train/"+str(m)+".ppm")
    #         gt = io.loadmat(i[1])
    #         gt = gt['LabelMap']
    #         ## print(i)
    #         # img = Image.open(i[1])
    #         # gt = file2array(i[1])
    #         # gt=np.loadtxt(i[1], dtype=np.int, delimiter=',')
    #
    #         # from numpy import genfromtxt
    #         # gt = genfromtxt(i[1], delimiter=',')
    #         img = Image.fromarray(gt.astype('uint8')).convert()
    #         img = img.resize((int(img.width*x),int(img.height*y)))
    #         # name = (i[1].split("/")[-1])
    #         # img.save("/home/hujin/SI/SI/trainset/rgbd/gt/"+str(m)+".png")
    #         if img.mode == "P":
    #             img = img.convert('RGB')
    #         img.save("/home/hujin/superpixel-benchmark/dataset/pascal/gt/"+str(m)+".png")
    #         m += 1


if __name__ == '__main__':
    import yaml

    with open('config.yml', 'r', encoding='utf-8') as f:
        cfg = yaml.load(f.read(), Loader=yaml.FullLoader)
    dataset = cfg["dataset"]
    x = cfg["scale_x"]
    y = cfg["scale_y"]

    if dataset not in ava_dataset:
        print("dataset name error")
        print(ava_dataset)
        exit()

    exp_name = cfg["exp_name"]
    try:
        x = float(x)
    except ValueError:
        print("value error")
        exit()
    try:
        y = float(y)
    except ValueError:
        print("value error")
        exit()
    if not (x > 0 and x <= 1):
        print("value error")
        exit()
    if not (y > 0 and y <= 1):
        print("value error")
        exit()
    init_dataset()

    for i in [100, 200, 300, 400, 500, 600]:
        n_seg = i # int(i * 1.05 + 1)
        test_name = exp_name + str(i)
        print(settings)
        print(test_name, file=log_file, flush=True)
        print(n_seg)
        process_dataset()
    # exit(0)
    # for i in [50]:
    #     for t in [0.25, 0.3, 0.4, 0.5, 0.6, .8, 1.0]:
    #         for diffu in [0.6, .8, 1.0, 1.2, 1.5, 2.0]:
    #             n_seg = int(i*1.05 + 1)
    #             test_name = exp_name + str(i)
    #             settings[2] = t
    #             settings[1] = diffu
    #             print(settings)
    #             print(test_name, file=log_file, flush=True)
    #             print(n_seg)
    #     # for t in [1/4, 1/3, 0.4, 0.5, 0.6, 0.8, 1.0, 1.2, 1.5]:
    #     #     settings[2] = t
    #     #     print(settings)
    #             process_dataset()
    # break

